#!/usr/bin/env python3
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation


def fitness_function(x1, x2):
    f1=x1-2*x2+3
    f2=2*x1+x2-8
    z = f1**2+f2**2
    return z

def update_velocity(particle, velocity, pbest, gbest, w_min=0.5, w_max=1.0, c=0.1):
    # Randomly generate r1, r2 and inertia weight from normal distribution
    r1 = np.random.uniform(0, w_max)
    r2 = np.random.uniform(0, w_max)
    w = np.random.uniform(w_min, w_max)
    c1 = c
    c2 = c
    new_velocity = w* velocity + c1*r1*(pbest - particle) + c2*r2*(gbest-particle)
    return new_velocity

def update_position(particle, velocity):
    # Move particles by adding velocity
    new_particle = particle + velocity
    return new_particle

def pso_2d(population, dimension, position_min, position_max, generation, fitness_criterion):
    history = []

    # Initialisation
    particles = np.random.uniform(position_min, position_max, size=(population, dimension))
    current_fitness = fitness_function(particles[:,0], particles[:,1])
    history = [[particles, current_fitness,],]

    # Particle's best position  &  fitness
    pbest_position = particles
    pbest_fitness = current_fitness

    # Index & position of the best particle
    gbest_index = np.argmin(pbest_fitness)
    gbest_position = pbest_position[gbest_index]

    # Velocity (starting from 0 speed)
    velocity = np.zeros([population, dimension])

    # Loop for the number of generation
    for t in range(generation):
        # Stop if the average fitness value reached a predefined success criterion
        if np.average(pbest_fitness) <= fitness_criterion:
            break

        velocity = update_velocity(particles, velocity, pbest_position, gbest_position)
        particles = update_position(particles, velocity)
        current_fitness = fitness_function(particles[:,0], particles[:,1])
        history.append([particles, current_fitness])

        # Update personal bests
        update_mask = current_fitness < pbest_fitness
        pbest_position[update_mask] = particles[update_mask]
        pbest_fitness[update_mask] = current_fitness[update_mask]

        # Update global best
        gbest_index = np.argmin(pbest_fitness)
        gbest_position = pbest_position[gbest_index]

    # Print the results
    print('Global Best Position: ', gbest_position)
    print('Best Fitness Value: ', min(pbest_fitness))
    print('Average Particle Best Fitness Value: ', np.average(pbest_fitness))
    print('Number of Generation: ', t)
    return history

def visualization(history, position_min, position_max):
    # Plotting prepartion
    fig = plt.figure(figsize=(10, 10))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    x = np.linspace(position_min, position_max, 80)
    y = np.linspace(position_min, position_max, 80)
    X, Y = np.meshgrid(x, y)
    Z = fitness_function(X,Y)
    ax.plot_wireframe(X, Y, Z, color='r', linewidth=0.2)

    scat = ax.scatter3D([], [], [], c='b')
    def _update(frame):
        particles, pbest_fitness = history[frame]
        scat._offsets3d = (particles[:, 0], particles[:, 1], pbest_fitness)
        return scat,

    _ani = FuncAnimation(fig, _update, frames=len(history), blit=False, repeat=True)
    plt.show()

if __name__ == '__main__':
    population = 100
    dimension = 2
    position_min = -100.0
    position_max = 100.0
    generation = 400
    fitness_criterion = 10e-4

    history = pso_2d(population, dimension, position_min, position_max, generation, fitness_criterion)
    visualization(history, position_min, position_max)

